import torch
import torch.nn as nn
import torch.nn.init as init


def _weights_init(m):
    classname = m.__class__.__name__
    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
        init.kaiming_normal_(m.weight)
    elif classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


class Generator(nn.Module):
    def __init__(self, opt):
        super(Generator, self).__init__()
        self.gru_type = opt.gru_type
        self.rnn = nn.GRU(input_size=opt.z_dim, hidden_size=opt.hidden_dim, num_layers=opt.num_layer,
                          batch_first=True)
        self.fc = nn.Linear(opt.hidden_dim, opt.z_dim)
        self.sigmoid = nn.Sigmoid()
        self.apply(_weights_init)

    def forward(self, input, delta_pre, sigmoid=True):
        g_outputs, _ = self.rnn(input)
        E = self.fc(g_outputs)
        if sigmoid:
            E = self.sigmoid(E)
        return E


class Supervisor(nn.Module):
    def __init__(self, opt):
        super(Supervisor, self).__init__()
        self.gru_type = opt.gru_type

        self.rnn = nn.GRU(input_size=opt.z_dim, hidden_size=opt.hidden_dim, num_layers=opt.num_layer,
                          batch_first=True)
        self.fc = nn.Linear(opt.hidden_dim, opt.z_dim)
        self.sigmoid = nn.Sigmoid()
        self.apply(_weights_init)

    def forward(self, input, delta_pre, sigmoid=True):
        s_outputs, _ = self.rnn(input)
        S = self.fc(s_outputs)
        if sigmoid:
            S = self.sigmoid(S)
        return S


class Discriminator(nn.Module):
    def __init__(self, opt):
        super(Discriminator, self).__init__()
        self.gru_type = opt.gru_type

        self.rnn = nn.GRU(input_size=opt.z_dim, hidden_size=opt.hidden_dim, num_layers=opt.num_layer,
                          batch_first=True)
        self.fc = nn.Linear(opt.hidden_dim, 1)
        self.sigmoid = nn.Sigmoid()
        self.apply(_weights_init)

    def forward(self, input, delta_pre, sigmoid=False):  # 这里的sigmoid，不要用，cross_entropy自带sigmoid, BCEWithSigmoid也带sigmoid
        d_outputs, _ = self.rnn(input)  # 这里用的是 output[bsz, seq_len, hidden_dim]，不是最后一个隐藏单元的状态，

        # d_outputs = d_outputs[:,-1,:] # todo 改为只要最有一步的输出，问题是，损失函数只有一部分了（之前有seq_len个），相当于减小了学习率
        Y_hat = self.fc(d_outputs)
        if sigmoid:
            Y_hat = self.sigmoid(Y_hat)
        return Y_hat


class Imputation(nn.Module):
    def __init__(self, opt, x):
        super(Imputation, self).__init__()

        self.z = torch.nn.Parameter(torch.tensor(x, dtype=torch.float32))

    def forward(self):
        return self.z
    # def __init__(self, opt):
    #     super(Imputation, self).__init__()
    #     self.z = torch.nn.Parameter(torch.randn(opt.batch_size, opt.seq_len, opt.z_dim))

    def forward(self):
        return self.z

